from PyQt5.QtWidgets import (QApplication, QMainWindow, QWidget, QVBoxLayout, 
                           QHBoxLayout, QLabel, QLineEdit, QPushButton, 
                           QGroupBox, QTextEdit, QMessageBox, QProgressBar, QComboBox,
                           QTabWidget, QGridLayout, QSpinBox)
from PyQt5.QtCore import Qt, QThread, pyqtSignal
from core.n5171b_control import N5171B
import sys
import numpy as np
from time import sleep
from core.data_collector import DataCollector
from config.config_manager import get_config_value
import time
from concurrent.futures import ThreadPoolExecutor, as_completed

class MultiChannelSweepThread(QThread):
    """多路信号幅值扫描的线程"""
    progress_updated = pyqtSignal(int, float, float)  # 发送进度、当前幅值和频率
    sweep_completed = pyqtSignal()  # 扫描完成信号
    data_saved = pyqtSignal(int, float, float)  # 数据保存完成信号，参数是通道号、当前幅值和频率
    status_updated = pyqtSignal(str)  # 添加状态更新信号
    
    def __init__(self, signal_generator, start_amp, stop_amp, step_amp, channel_count):
        super().__init__()
        self.signal_generator = signal_generator
        self.start_amp = start_amp
        self.stop_amp = stop_amp
        self.step_amp = step_amp
        self.channel_count = channel_count
        self.is_running = True
        self.data_collectors = [None] * channel_count  # 每个通道一个数据采集器
        self.thread_pool = ThreadPoolExecutor(max_workers=channel_count)  # 创建线程池

    def run(self):
        # 计算所有幅值点
        amplitudes = np.arange(self.start_amp, self.stop_amp + self.step_amp, self.step_amp)
        total_points = len(amplitudes)
        
        for i, amp in enumerate(amplitudes):
            if not self.is_running:
                break
                
            # 设置新的幅值
            self.status_updated.emit(f"正在设置幅值: {amp:.2f} dBm...")
            self.signal_generator.set_amplitude(amp)
            
            # 等待设备稳定，从配置读取延迟时间
            stabilization_time = get_config_value("data_collection.stabilization_time", 2.0)
            self.status_updated.emit(f"等待设备稳定 ({stabilization_time}s)...")
            sleep(stabilization_time)
            
            # 获取当前频率
            current_freq = self.signal_generator.get_frequency()
            
            # 并行采集所有通道的数据
            futures = []
            for channel in range(self.channel_count):
                if not self.is_running:
                    break
                self.status_updated.emit(f"启动通道 {channel+1} 数据采集 ({amp:.2f} dBm)...")
                future = self.thread_pool.submit(self.collect_data, channel, amp, current_freq)
                futures.append((channel, future))
            
            # 等待所有通道采集完成
            for channel, future in futures:
                try:
                    success = future.result()
                    if not success:
                        self.status_updated.emit(f"通道 {channel+1} 在幅值 {amp:.2f} dBm 时数据采集失败")
                except Exception as e:
                    self.status_updated.emit(f"通道 {channel+1} 数据采集异常: {str(e)}")
            
            # 发送进度更新
            progress = int((i + 1) / total_points * 100)
            self.progress_updated.emit(progress, amp, current_freq)
        
        # 关闭线程池
        self.thread_pool.shutdown(wait=True)
        self.sweep_completed.emit()

    def collect_data(self, channel, current_amp, current_freq):
        """数据采集函数"""
        try:
            collector = self.data_collectors[channel]
            if not collector:
                return False
                
            self.status_updated.emit(f"  - 通道 {channel+1} 正在采集中...")
            success = collector.save_data(current_amp, current_freq)
            if not success:
                return False
                
            self.data_saved.emit(channel, current_amp, current_freq)
            return True
            
        except Exception as e:
            self.status_updated.emit(f"通道 {channel+1} 数据采集错误: {str(e)}")
            return False

    def stop(self):
        self.is_running = False
        if hasattr(self, 'thread_pool'):
            self.thread_pool.shutdown(wait=False)

class ChannelTab(QWidget):
    """单个通道的控制标签页"""
    def __init__(self, channel_num):
        super().__init__()
        self.channel_num = channel_num
        self.setup_ui()
        
    def setup_ui(self):
        layout = QVBoxLayout()
        
        # 通道选择
        channel_group = QGroupBox(f"通道 {self.channel_num} 设置")
        channel_layout = QGridLayout()
        
        # 通道选择下拉框
        channel_layout.addWidget(QLabel("数据采集通道:"), 0, 0)
        self.channel_combo = QComboBox()
        self.channel_combo.addItems(get_config_value("epics.channels", {}).keys())
        channel_layout.addWidget(self.channel_combo, 0, 1)
        
        # 测试按钮
        self.test_btn = QPushButton("测试通道")
        channel_layout.addWidget(self.test_btn, 0, 2)
        
        # 状态显示
        self.status_label = QLabel("未连接")
        channel_layout.addWidget(self.status_label, 1, 0, 1, 3)
        
        channel_group.setLayout(channel_layout)
        layout.addWidget(channel_group)
        
        # 数据显示区域
        display_group = QGroupBox("数据显示")
        display_layout = QVBoxLayout()
        self.display_text = QTextEdit()
        self.display_text.setReadOnly(True)
        self.display_text.setFixedHeight(100)
        display_layout.addWidget(self.display_text)
        display_group.setLayout(display_layout)
        layout.addWidget(display_group)
        
        self.setLayout(layout)

class MultiChannelGUI(QMainWindow):
    def __init__(self):
        super().__init__()
        self.setWindowTitle("多路信号源控制(By: Xu)")
        self.setFixedSize(800, 900)
        
        # 信号源实例
        self.signal_generator = None
        self.sweep_thread = None
        
        # 保存当前值的变量
        self.current_amplitude = 0
        self.current_frequency = 0
        
        # 创建主窗口部件和布局
        main_widget = QWidget()
        self.setCentralWidget(main_widget)
        layout = QVBoxLayout(main_widget)
        
        # 创建各个区域
        self.create_connection_group(layout)
        self.create_control_group(layout)
        self.create_channel_tabs(layout)
        self.create_sweep_group(layout)
        self.create_button_group(layout)
        
        # 设置样式
        self.apply_styles()

    def create_connection_group(self, parent_layout):
        group = QGroupBox("设备连接")
        layout = QHBoxLayout()
        
        # IP地址输入
        layout.addWidget(QLabel("IP地址:"))
        self.ip_edit = QLineEdit(get_config_value("signal_generator.ip_address", "127.0.0.1"))
        layout.addWidget(self.ip_edit)
        
        # 连接按钮
        self.connect_btn = QPushButton("连接")
        self.connect_btn.clicked.connect(self.connect_device)
        layout.addWidget(self.connect_btn)
        
        # 状态显示
        self.status_label = QLabel("未连接")
        layout.addWidget(self.status_label)
        
        group.setLayout(layout)
        parent_layout.addWidget(group)

    def create_control_group(self, parent_layout):
        group = QGroupBox("参数控制")
        layout = QGridLayout()
        
        # 幅值控制
        layout.addWidget(QLabel("幅值 (dBm):"), 0, 0)
        self.amp_edit = QLineEdit()
        self.amp_edit.setPlaceholderText("输入幅值")
        layout.addWidget(self.amp_edit, 0, 1)
        
        # 频率控制
        layout.addWidget(QLabel("频率 (MHz):"), 1, 0)
        self.freq_edit = QLineEdit()
        self.freq_edit.setPlaceholderText("输入频率")
        layout.addWidget(self.freq_edit, 1, 1)
        
        # 通道数量控制
        layout.addWidget(QLabel("通道数量:"), 2, 0)
        self.channel_count_spin = QSpinBox()
        self.channel_count_spin.setRange(1, 8)  # 最多支持8个通道
        self.channel_count_spin.setValue(4)  # 默认4个通道
        self.channel_count_spin.valueChanged.connect(self.update_channel_tabs)
        layout.addWidget(self.channel_count_spin, 2, 1)
        
        group.setLayout(layout)
        parent_layout.addWidget(group)

    def create_channel_tabs(self, parent_layout):
        self.tab_widget = QTabWidget()
        self.channel_tabs = []
        
        # 默认创建4个通道标签页
        for i in range(4):
            tab = ChannelTab(i + 1)
            tab.test_btn.clicked.connect(lambda checked, num=i: self.test_channel(num))
            self.channel_tabs.append(tab)
            self.tab_widget.addTab(tab, f"通道 {i + 1}")
        
        parent_layout.addWidget(self.tab_widget)

    def create_sweep_group(self, parent_layout):
        group = QGroupBox("幅值扫描控制")
        layout = QGridLayout()
        
        # 起始幅值
        layout.addWidget(QLabel("起始幅值 (dBm):"), 0, 0)
        self.start_amp_edit = QLineEdit("-30")
        layout.addWidget(self.start_amp_edit, 0, 1)
        
        # 终止幅值
        layout.addWidget(QLabel("终止幅值 (dBm):"), 1, 0)
        self.stop_amp_edit = QLineEdit("10")
        layout.addWidget(self.stop_amp_edit, 1, 1)
        
        # 步进值
        layout.addWidget(QLabel("步进值 (dB):"), 2, 0)
        self.step_amp_edit = QLineEdit("0.5")
        layout.addWidget(self.step_amp_edit, 2, 1)
        
        # 进度条
        self.progress_bar = QProgressBar()
        self.progress_bar.setAlignment(Qt.AlignCenter)
        layout.addWidget(self.progress_bar, 3, 0, 1, 2)
        
        # 扫描控制按钮
        button_layout = QHBoxLayout()
        self.start_sweep_btn = QPushButton("开始扫描")
        self.start_sweep_btn.clicked.connect(self.start_sweep)
        self.stop_sweep_btn = QPushButton("停止扫描")
        self.stop_sweep_btn.clicked.connect(self.stop_sweep)
        self.stop_sweep_btn.setEnabled(False)
        button_layout.addWidget(self.start_sweep_btn)
        button_layout.addWidget(self.stop_sweep_btn)
        layout.addLayout(button_layout, 4, 0, 1, 2)
        
        group.setLayout(layout)
        parent_layout.addWidget(group)

    def create_button_group(self, parent_layout):
        layout = QHBoxLayout()
        
        # 读取按钮
        read_btn = QPushButton("读取当前值")
        read_btn.clicked.connect(self.read_values)
        layout.addWidget(read_btn)
        
        # 设置按钮
        set_btn = QPushButton("设置新值")
        set_btn.clicked.connect(self.set_values)
        layout.addWidget(set_btn)
        
        # 断开连接按钮
        disconnect_btn = QPushButton("断开连接")
        disconnect_btn.clicked.connect(self.disconnect_device)
        layout.addWidget(disconnect_btn)
        
        parent_layout.addLayout(layout)

    def apply_styles(self):
        self.setStyleSheet("""
            QMainWindow {
                background-color: #f0f0f0;
            }
            QGroupBox {
                font-weight: bold;
                border: 2px solid #cccccc;
                border-radius: 6px;
                margin-top: 6px;
            }
            QGroupBox::title {
                subcontrol-origin: margin;
                left: 7px;
                padding: 0px 5px 0px 5px;
            }
            QPushButton {
                background-color: #2196F3;
                color: white;
                border: none;
                padding: 5px 15px;
                border-radius: 4px;
                min-width: 80px;
            }
            QPushButton:hover {
                background-color: #1976D2;
            }
            QPushButton:pressed {
                background-color: #0D47A1;
            }
            QPushButton:disabled {
                background-color: #cccccc;
            }
            QLineEdit {
                padding: 5px;
                border: 1px solid #cccccc;
                border-radius: 4px;
            }
            QTextEdit {
                border: 1px solid #cccccc;
                border-radius: 4px;
            }
        """)

    def update_channel_tabs(self):
        count = self.channel_count_spin.value()
        current_count = len(self.channel_tabs)
        
        if count > current_count:
            # 添加新的标签页
            for i in range(current_count, count):
                tab = ChannelTab(i + 1)
                tab.test_btn.clicked.connect(lambda checked, num=i: self.test_channel(num))
                self.channel_tabs.append(tab)
                self.tab_widget.addTab(tab, f"通道 {i + 1}")
        elif count < current_count:
            # 移除多余的标签页
            for i in range(current_count - 1, count - 1, -1):
                self.tab_widget.removeTab(i)
                self.channel_tabs.pop()

    def connect_device(self):
        try:
            ip = self.ip_edit.text()
            self.signal_generator = N5171B(ip)
            self.status_label.setText("已连接")
            self.status_label.setStyleSheet("color: green; font-weight: bold;")
            self.read_values()
        except Exception as e:
            QMessageBox.critical(self, "错误", f"连接失败: {str(e)}")

    def disconnect_device(self):
        if self.signal_generator:
            self.signal_generator.close()
            self.signal_generator = None
            self.status_label.setText("未连接")
            self.status_label.setStyleSheet("")
            for tab in self.channel_tabs:
                tab.display_text.clear()

    def read_values(self):
        if not self.signal_generator:
            QMessageBox.warning(self, "警告", "请先连接设备")
            return
            
        try:
            amplitude = self.signal_generator.get_amplitude()
            frequency = self.signal_generator.get_frequency()
            
            # 保存当前值
            self.current_amplitude = amplitude
            self.current_frequency = frequency
            
            # 更新显示
            for tab in self.channel_tabs:
                tab.display_text.setPlainText(
                    f"当前设置：\n"
                    f"幅值: {amplitude:.2f} dBm\n"
                    f"频率: {frequency/1e6:.2f} MHz"
                )
        except Exception as e:
            QMessageBox.critical(self, "错误", f"读取失败: {str(e)}")

    def set_values(self):
        """设置幅值和频率"""
        if not self.signal_generator:
            QMessageBox.warning(self, "错误", "请先连接设备")
            return
            
        try:
            # 设置频率
            freq_mhz = float(self.freq_edit.text())
            self.signal_generator.set_frequency(freq_mhz * 1e6)
            
            # 手动设置幅值的功能在此简化或移除，由扫描控制
            QMessageBox.information(self, "成功", "频率设置成功")
            
            self.read_values()
        except ValueError:
            QMessageBox.warning(self, "错误", "无效的频率值")
        except Exception as e:
            QMessageBox.critical(self, "错误", f"设置失败: {str(e)}")

    def test_channel(self, channel_num):
        """测试指定通道的连接"""
        tab = self.channel_tabs[channel_num]
        selected_channel = tab.channel_combo.currentText()
        
        # 禁用测试按钮，显示正在测试
        tab.test_btn.setEnabled(False)
        tab.test_btn.setText("测试中...")
        QApplication.processEvents()  # 确保UI更新
        
        try:
            data_collector = DataCollector()
            success, message = data_collector.test_channel(selected_channel)
            if success:
                tab.status_label.setText(f"通道测试成功: {message}")
                tab.status_label.setStyleSheet("color: green;")
            else:
                tab.status_label.setText(f"通道测试失败: {message}")
                tab.status_label.setStyleSheet("color: red;")
        except Exception as e:
            tab.status_label.setText(f"测试错误: {str(e)}")
            tab.status_label.setStyleSheet("color: red;")
        finally:
            # 恢复测试按钮
            tab.test_btn.setEnabled(True)
            tab.test_btn.setText("测试通道")

    def start_sweep(self):
        """开始幅值扫描"""
        if not self.signal_generator:
            QMessageBox.warning(self, "警告", "请先连接设备")
            return
            
        try:
            # 获取通道数量
            channel_count = self.channel_count_spin.value()
            
            # 检查每个通道的设置
            data_collectors = []
            for i in range(channel_count):
                tab = self.channel_tabs[i]
                selected_channel = tab.channel_combo.currentText()
                if not selected_channel:
                    QMessageBox.warning(self, "警告", f"请为通道 {i+1} 选择数据采集通道")
                    return
                
                # 创建数据采集器
                data_collector = DataCollector()
                success, message = data_collector.test_channel(selected_channel)
                if not success:
                    QMessageBox.critical(self, "错误", f"通道 {i+1} 测试失败: {message}")
                    return
                
                # 连接通道并创建新的数据文件
                data_collector.connect_channel(selected_channel)
                data_collector.start_new_measurement()
                data_collectors.append(data_collector)
            
            # 获取扫描参数
            start_amp = float(self.start_amp_edit.text())
            stop_amp = float(self.stop_amp_edit.text())
            step_amp = float(self.step_amp_edit.text())
            
            # 检查输入值的合理性
            if start_amp > stop_amp:
                QMessageBox.warning(self, "警告", "起始幅值不能大于终止幅值")
                return
            if step_amp <= 0:
                QMessageBox.warning(self, "警告", "步进值必须大于0")
                return
                
            # 创建并启动扫描线程
            self.sweep_thread = MultiChannelSweepThread(
                self.signal_generator, start_amp, stop_amp, step_amp, channel_count)
            self.sweep_thread.data_collectors = data_collectors
            
            self.sweep_thread.progress_updated.connect(self.update_sweep_progress)
            self.sweep_thread.sweep_completed.connect(self.sweep_completed)
            self.sweep_thread.data_saved.connect(self.on_data_saved)
            self.sweep_thread.status_updated.connect(self.update_status)
            
            # 启动线程
            self.sweep_thread.start()
            
            # 更新按钮状态
            self.start_sweep_btn.setEnabled(False)
            self.stop_sweep_btn.setEnabled(True)
            
            # 清空显示
            for tab in self.channel_tabs:
                tab.display_text.clear()
            
        except ValueError:
            QMessageBox.critical(self, "错误", "请输入有效的数值")
        except Exception as e:
            QMessageBox.critical(self, "错误", f"启动扫描失败: {str(e)}")

    def stop_sweep(self):
        """停止幅值扫描"""
        if self.sweep_thread and self.sweep_thread.isRunning():
            self.sweep_thread.stop()
            self.sweep_thread.wait()
            self.sweep_completed()

    def update_sweep_progress(self, progress, current_amp, current_freq):
        """更新扫描进度"""
        self.progress_bar.setValue(progress)
        status_text = f"当前：幅值 = {current_amp:.2f} dBm, 频率 = {current_freq/1e6:.2f} MHz"
        for tab in self.channel_tabs:
            tab.display_text.append(status_text)

    def sweep_completed(self):
        """扫描完成处理"""
        self.start_sweep_btn.setEnabled(True)
        self.stop_sweep_btn.setEnabled(False)
        self.progress_bar.setValue(0)
        
        # 关闭所有数据文件
        if hasattr(self, 'sweep_thread') and self.sweep_thread.data_collectors:
            for collector in self.sweep_thread.data_collectors:
                if collector:
                    collector.close()
        
        for tab in self.channel_tabs:
            tab.display_text.append("扫描完成")

    def on_data_saved(self, channel, amp, freq):
        """处理数据保存完成信号"""
        if channel < len(self.channel_tabs):
            tab = self.channel_tabs[channel]
            collector = self.sweep_thread.data_collectors[channel]
            tab.display_text.append(
                f">>> 数据存储完成 [{collector.current_channel}] "
                f"(第{collector.collect_count}次): "
                f"幅值 = {amp:.2f} dBm, 频率 = {freq/1e6:.2f} MHz")

    def update_status(self, message):
        """更新状态显示"""
        for tab in self.channel_tabs:
            tab.display_text.append(message)

def main():
    app = QApplication(sys.argv)
    window = MultiChannelGUI()
    window.show()
    sys.exit(app.exec_())

if __name__ == "__main__":
    main()